{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用 tf.funcation 的性能\n",
    "tensorflow2 集合了 eager 模式的便利性和 tensorflow1 的强大功能。合并的核心是  tf.function，它允许我们将 python 语法转换为可移植的、高性能的 tensorflow 图。\n",
    "\n",
    "tf.function 一个很酷的新功能是 AutoGraph，它使你可以使用 python 语法写图形代码。有关可与AutoGraph一起使用的Python功能的列表，请参见 [AutoGraph Capabilities and Limitation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md)。有关tf.function的更多详细信息，请参阅 [TF 2.0: Functions, not Sessions](https://github.com/tensorflow/community/blob/master/rfcs/20180918-functions-not-sessions-20.md)。有关AutoGraph的更多详细信息，请参见 [tf.autograph](https://tensorflow.google.cn/api_docs/python/tf/autograph)。\n",
    "\n",
    "这个教程将向你介绍 tf.function 和 AutoGraph 的基本功能"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 起步\n",
    "导入 tensorflow2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 tf.function 装饰\n",
    "当你在一个函数上注解了 tf.function，你仍然可以像调用其他普通的 python 函数一样调用它，但是它将会被编译成图表，这意味着你可以受益于更快的执行速度，在 GPU 或 CPU 上运行或导出到 SavedModel。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3, 3), dtype=float32, numpy=\n",
       "array([[0.44946098, 1.003785  , 0.39791206],\n",
       "       [0.60381234, 0.9506065 , 0.5102525 ],\n",
       "       [0.63267857, 0.67906773, 0.38511965]], dtype=float32)>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@tf.function\n",
    "def simple_nn_layer(x, y):\n",
    "    return tf.nn.relu(tf.matmul(x, y))\n",
    "\n",
    "x = tf.random.uniform((3, 3))\n",
    "y = tf.random.uniform((3, 3))\n",
    "\n",
    "simple_nn_layer(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果我们检查注释的结果，我们会发现它是一个特殊的可调用对象，可处理与TensorFlow运行时的所有交互。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.eager.def_function.Function at 0x1786c169648>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "simple_nn_layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果您的代码使用多个函数，则无需全部注释——从带注释的函数调用的任何函数也将在图形模式下运行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(3,), dtype=int32, numpy=array([3, 5, 7])>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def linear_layer(x):\n",
    "  return 2 * x + 1\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def deep_net(x):\n",
    "  return tf.nn.relu(linear_layer(x))\n",
    "\n",
    "\n",
    "deep_net(tf.constant((1, 2, 3)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于具有许多小操作的图形，功能可能比 eager 的代码更快。但是对于具有一些耗性能的操作（例如卷积）的图，你可能看不到太多的加速。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eager conv: 0.4032915000002504\n",
      "Function conv: 0.39127019999978074\n",
      "eager lstm: 0.005612899999960064\n",
      "function lstm: 0.0039416999998138635\n"
     ]
    }
   ],
   "source": [
    "import timeit\n",
    "conv_layer = tf.keras.layers.Conv2D(100, 3)\n",
    "\n",
    "@tf.function\n",
    "def conv_fn(image):\n",
    "  return conv_layer(image)\n",
    "\n",
    "image = tf.zeros([1, 200, 200, 100])\n",
    "conv_layer(image); conv_fn(image)\n",
    "print(\"Eager conv:\", timeit.timeit(lambda: conv_layer(image), number=10))\n",
    "print(\"Function conv:\", timeit.timeit(lambda: conv_fn(image), number=10))\n",
    "\n",
    "lstm_cell = tf.keras.layers.LSTMCell(10)\n",
    "\n",
    "@tf.function\n",
    "def lstm_fn(input, state):\n",
    "  return lstm_cell(input, state)\n",
    "\n",
    "input = tf.zeros([10, 10])\n",
    "state = [tf.zeros([10, 10])] * 2\n",
    "# warm up\n",
    "lstm_cell(input, state); lstm_fn(input, state)\n",
    "print(\"eager lstm:\", timeit.timeit(lambda: lstm_cell(input, state), number=10))\n",
    "print(\"function lstm:\", timeit.timeit(lambda: lstm_fn(input, state), number=10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用Python控制流程\n",
    "在 tf.function 在使用依赖数据的控制流时，你能使用 python 控制流语句，AutoGraph 会将其转换为合适的 tensorflow 操作。比如，如果语句依赖于张量，则将其转换为 tf.cond()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "square_if_positive(2) = 4\n",
      "square_if_positive(-2) = 0\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def square_if_positive(x):\n",
    "  if x > 0:\n",
    "    x = x * x\n",
    "  else:\n",
    "    x = 0\n",
    "  return x\n",
    "\n",
    "\n",
    "print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))\n",
    "print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AutoGraph支持常见的 python 语句，例如 while、for、if、break、continue 和 return，并支持嵌套。这意味着你可以在 while 和 if 语句的条件下使用 tensorflow 表达式，或者在 for 循环中遍历 tensor。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=int32, numpy=42>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@tf.function\n",
    "def sum_even(items):\n",
    "  s = 0\n",
    "  for c in items:\n",
    "    if c % 2 > 0:\n",
    "      continue\n",
    "    s += c\n",
    "  return s\n",
    "\n",
    "\n",
    "sum_even(tf.constant([10, 12, 15, 20]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AutoGraph 还为高级用户提供了一个底层 API。例如，我们可以使用它来查看生成的代码。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def tf__sum_even(items):\n",
      "  do_return = False\n",
      "  retval_ = ag__.UndefinedReturnValue()\n",
      "  with ag__.FunctionScope('sum_even', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:\n",
      "    s = 0\n",
      "\n",
      "    def get_state_2():\n",
      "      return ()\n",
      "\n",
      "    def set_state_2(_):\n",
      "      pass\n",
      "\n",
      "    def loop_body(iterates, s):\n",
      "      c = iterates\n",
      "      continue_ = False\n",
      "\n",
      "      def get_state():\n",
      "        return ()\n",
      "\n",
      "      def set_state(_):\n",
      "        pass\n",
      "\n",
      "      def if_true():\n",
      "        continue_ = True\n",
      "        return continue_\n",
      "\n",
      "      def if_false():\n",
      "        return continue_\n",
      "      cond = c % 2 > 0\n",
      "      continue_ = ag__.if_stmt(cond, if_true, if_false, get_state, set_state, ('continue_',), ())\n",
      "\n",
      "      def get_state_1():\n",
      "        return ()\n",
      "\n",
      "      def set_state_1(_):\n",
      "        pass\n",
      "\n",
      "      def if_true_1():\n",
      "        s_1, = s,\n",
      "        s_1 += c\n",
      "        return s_1\n",
      "\n",
      "      def if_false_1():\n",
      "        return s\n",
      "      cond_1 = ag__.not_(continue_)\n",
      "      s = ag__.if_stmt(cond_1, if_true_1, if_false_1, get_state_1, set_state_1, ('s',), ())\n",
      "      return s,\n",
      "    s, = ag__.for_stmt(items, None, loop_body, get_state_2, set_state_2, (s,), ('s',), ())\n",
      "    do_return = True\n",
      "    retval_ = fscope.mark_return_value(s)\n",
      "  do_return,\n",
      "  return ag__.retval(retval_)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(tf.autograph.to_code(sum_even.python_function))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keras 和 AutoGraph\n",
    "默认情况下，AutoGraph 在非动态 Keras 模型中可用。有关更多信息，可参见 [tf.keras](https://tensorflow.google.cn/api_docs/python/tf/keras)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=int32, numpy=array([-1, -2])>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class CustomModel(tf.keras.models.Model):\n",
    "\n",
    "  @tf.function\n",
    "  def call(self, input_data):\n",
    "    if tf.reduce_mean(input_data) > 0:\n",
    "      return input_data\n",
    "    else:\n",
    "      return input_data // 2\n",
    "\n",
    "\n",
    "model = CustomModel()\n",
    "\n",
    "model(tf.constant([-2, -4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 额外效果\n",
    "就像在 eager 模式下一样，你可以使用带有额外效果的操作，例如通常在 tf.function 中使用 tf.assign 或tf.print，并且它将插入必要的控件依赖项以确保它们按顺序执行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'Variable:0' shape=() dtype=int32, numpy=7>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v = tf.Variable(5)\n",
    "\n",
    "@tf.function\n",
    "def find_next_odd():\n",
    "  v.assign(v + 1)\n",
    "  if v % 2 == 0:\n",
    "    v.assign(v + 1)\n",
    "\n",
    "\n",
    "find_next_odd()\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 调试\n",
    "tf.function 和 AutoGraph 通过生成代码并将其跟踪到 tensorflow 图中来工作。该机制尚不支持像 pdb 的逐步调试器。但是，你可以调用 tf.config.experimental_run_functions_eagerly（True）来临时启用   \"tf.function\" 内部的急切执行并使用您喜欢的调试器："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "  if x > 0:\n",
    "    x = x + 1\n",
    "  return x\n",
    "\n",
    "tf.config.experimental_run_functions_eagerly(True)\n",
    "\n",
    "f(tf.constant(1))\n",
    "\n",
    "tf.config.experimental_run_functions_eagerly(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 高级例子：一个图内循环训练\n",
    "上一节显示了可以在Keras图层和模型中使用AutoGraph。 Keras模型也可以在 AutoGraph 代码中使用。 \n",
    "\n",
    "此示例显示了如何在图上执行在 MNIST 上训练整个训练过程（加载批次、计算梯度、更新参数、计算验证精度以及重复直到收敛）的简单Keras模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n",
      "11493376/11490434 [==============================] - 116s 10us/step\n",
      "Step 10 : loss 1.71716428 ; accuracy 0.364\n",
      "Step 20 : loss 1.16976023 ; accuracy 0.5055\n",
      "Step 30 : loss 0.616642475 ; accuracy 0.598333359\n",
      "Step 40 : loss 0.646155119 ; accuracy 0.6545\n",
      "Step 50 : loss 0.372862428 ; accuracy 0.6936\n",
      "Step 60 : loss 0.455450892 ; accuracy 0.722166657\n",
      "Step 70 : loss 0.291360289 ; accuracy 0.74585712\n",
      "Step 80 : loss 0.530088902 ; accuracy 0.764\n",
      "Step 90 : loss 0.289816976 ; accuracy 0.779\n",
      "Step 100 : loss 0.331202745 ; accuracy 0.7921\n",
      "Step 110 : loss 0.314658731 ; accuracy 0.802727282\n",
      "Step 120 : loss 0.281935781 ; accuracy 0.810833335\n",
      "Step 130 : loss 0.446094751 ; accuracy 0.818692327\n",
      "Step 140 : loss 0.267202318 ; accuracy 0.825928569\n",
      "Step 150 : loss 0.21517624 ; accuracy 0.831466675\n",
      "Step 160 : loss 0.215192378 ; accuracy 0.83743751\n",
      "Step 170 : loss 0.258279681 ; accuracy 0.841529429\n",
      "Step 180 : loss 0.254997343 ; accuracy 0.845388889\n",
      "Step 190 : loss 0.19334276 ; accuracy 0.850263178\n",
      "Step 200 : loss 0.308558434 ; accuracy 0.854\n",
      "Final step tf.Tensor(200, shape=(), dtype=int32) : loss tf.Tensor(0.30855843, shape=(), dtype=float32) ; accuracy tf.Tensor(0.854, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# 1. 下载数据集\n",
    "def prepare_mnist_features_and_labels(x, y):\n",
    "  x = tf.cast(x, tf.float32) / 255.0\n",
    "  y = tf.cast(y, tf.int64)\n",
    "  return x, y\n",
    "\n",
    "def mnist_dataset():\n",
    "  (x, y), _ = tf.keras.datasets.mnist.load_data()\n",
    "  ds = tf.data.Dataset.from_tensor_slices((x, y))\n",
    "  ds = ds.map(prepare_mnist_features_and_labels)\n",
    "  ds = ds.take(20000).shuffle(20000).batch(100)\n",
    "  return ds\n",
    "\n",
    "train_dataset = mnist_dataset()\n",
    "\n",
    "# 2. 定义模型\n",
    "model = tf.keras.Sequential((\n",
    "    tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),\n",
    "    tf.keras.layers.Dense(100, activation='relu'),\n",
    "    tf.keras.layers.Dense(100, activation='relu'),\n",
    "    tf.keras.layers.Dense(10)))\n",
    "model.build()\n",
    "optimizer = tf.keras.optimizers.Adam()\n",
    "\n",
    "# 3. 定义训练循环\n",
    "compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n",
    "\n",
    "\n",
    "def train_one_step(model, optimizer, x, y):\n",
    "  with tf.GradientTape() as tape:\n",
    "    logits = model(x)\n",
    "    loss = compute_loss(y, logits)\n",
    "\n",
    "  grads = tape.gradient(loss, model.trainable_variables)\n",
    "  optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
    "\n",
    "  compute_accuracy(y, logits)\n",
    "  return loss\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def train(model, optimizer):\n",
    "  train_ds = mnist_dataset()\n",
    "  step = 0\n",
    "  loss = 0.0\n",
    "  accuracy = 0.0\n",
    "  for x, y in train_ds:\n",
    "    step += 1\n",
    "    loss = train_one_step(model, optimizer, x, y)\n",
    "    if step % 10 == 0:\n",
    "      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())\n",
    "  return step, loss, accuracy\n",
    "\n",
    "step, loss, accuracy = train(model, optimizer)\n",
    "print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 批次\n",
    "在实际应用中，批处理对于性能至关重要。转换为 AutoGraph 的最佳代码是在批处理级别确定控制流的代码。如果要在各个示例级别做出决策，请尝试使用批处理API来保持性能。比如，如果你在 python 中有以下代码："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-5, -4, -3, -2, -1, 0, 1, 4, 9, 16]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def square_if_positive(x):\n",
    "  return [i ** 2 if i > 0 else i for i in x]\n",
    "\n",
    "\n",
    "square_if_positive(range(-5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "也许在 tensorflow 中你会这样写："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1,  0,  1,  4,  9, 16])>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@tf.function\n",
    "def square_if_positive_naive(x):\n",
    "  result = tf.TensorArray(tf.int32, size=x.shape[0])\n",
    "  for i in tf.range(x.shape[0]):\n",
    "    if x[i] > 0:\n",
    "      result = result.write(i, x[i] ** 2)\n",
    "    else:\n",
    "      result = result.write(i, x[i])\n",
    "  return result.stack()\n",
    "\n",
    "\n",
    "square_if_positive_naive(tf.range(-5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但是对于这个例子，你可以这样写："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1,  0,  1,  4,  9, 16])>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def square_if_positive_vectorized(x):\n",
    "  return tf.where(x > 0, x ** 2, x)\n",
    "\n",
    "\n",
    "square_if_positive_vectorized(tf.range(-5, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 再追踪\n",
    "主要点：\n",
    "1. 要谨慎调用带有非张量参数或具有变化形状的参数的函数\n",
    "2. 装饰模块级函数和模块级类的方法，避免修饰局部函数或方法。\n",
    "\n",
    "tf.function可以大大提高 eager 执行的速度，但会降低首次执行的速度。这是因为第一次执行该函数时，还会将该函数追踪到 tensorflow 图中。与实际执行图形相比，构造和优化图形通常要慢得多："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First invocation: 0.0595394000001761\n",
      "Second invocation: 0.0021884999996473198\n"
     ]
    }
   ],
   "source": [
    "import timeit\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def f(x, y):\n",
    "  return tf.matmul(x, y)\n",
    "\n",
    "print(\n",
    "    \"First invocation:\",\n",
    "    timeit.timeit(lambda: f(tf.ones((10, 10)), tf.ones((10, 10))), number=1))\n",
    "\n",
    "print(\n",
    "    \"Second invocation:\",\n",
    "    timeit.timeit(lambda: f(tf.ones((10, 10)), tf.ones((10, 10))), number=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以通过在函数顶部添加打印语句来轻松知道何时跟踪函数。因为任何Python代码仅在跟踪时执行，所以仅在跟踪函数时才会看到print的输出："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First invocation:\n",
      "Tracing!\n",
      "Executing\n",
      "Second invocation:\n",
      "Executing\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def f():\n",
    "  print('Tracing!')\n",
    "  tf.print('Executing')\n",
    "\n",
    "print('First invocation:')\n",
    "f()\n",
    "\n",
    "print('Second invocation:')\n",
    "f()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当使用不同的非张量参数调用时，tf.function也可能会重新跟踪："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 Tracing!\n",
      "1 Executing\n",
      "1 Executing\n",
      "2 Tracing!\n",
      "2 Executing\n",
      "2 Executing\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def f(n):\n",
    "  print(n, 'Tracing!')\n",
    "  tf.print(n, 'Executing')\n",
    "\n",
    "f(1)\n",
    "f(1)\n",
    "\n",
    "f(2)\n",
    "f(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当张量参数改变形状时也可能发生重新跟踪，除非指定了input_signature："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1,) Tracing!\n",
      "[1] Executing\n",
      "[2] Executing\n",
      "(2,) Tracing!\n",
      "[1 2] Executing\n",
      "[3 4] Executing\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "  print(x.shape, 'Tracing!')\n",
    "  tf.print(x, 'Executing')\n",
    "\n",
    "f(tf.constant([1]))\n",
    "f(tf.constant([2]))\n",
    "\n",
    "f(tf.constant([1, 2]))\n",
    "f(tf.constant([3, 4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此外，在调用 tf.function 时，它始终会使用其自己的一组轨迹创建一个新的graph函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracing!\n",
      "Executing\n",
      "Tracing!\n",
      "Executing\n"
     ]
    }
   ],
   "source": [
    "def f():\n",
    "  print('Tracing!')\n",
    "  tf.print('Executing')\n",
    "\n",
    "tf.function(f)()\n",
    "tf.function(f)()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在嵌套函数中使用 @tf.function装饰器时，这可能导致出乎意料的行为："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tracing!\n",
      "Executing\n",
      "Tracing!\n",
      "Executing\n"
     ]
    }
   ],
   "source": [
    "def outer():\n",
    "  @tf.function\n",
    "  def f():\n",
    "    print('Tracing!')\n",
    "    tf.print('Executing')\n",
    "  f()\n",
    "\n",
    "outer()\n",
    "outer()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
